#!/usr/bin/env python3
"""
ChromaDB 编码问题修复脚本

这个脚本可以帮助诊断和修复 ChromaDB 中因非标准 UTF-8 字符导致的编码问题。
"""

import sys
import os
from pathlib import Path

# 添加项目路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

try:
    from chromadb import HttpClient
    import orjson
    import unicodedata
    import numpy as np
except ImportError as e:
    print(f"缺少必要的依赖包: {e}")
    print("请安装依赖: pip install chromadb orjson numpy")
    sys.exit(1)

def clean_text(text):
    """
    清理文本中的非标准字符
    """
    if not isinstance(text, str):
        text = str(text)
    
    # 规范化 Unicode 字符
    text = unicodedata.normalize('NFKD', text)
    
    # 移除或替换代理对字符
    cleaned_text = ""
    for char in text:
        if not (0xD800 <= ord(char) <= 0xDFFF):  # 不是代理对字符
            cleaned_text += char
        else:
            cleaned_text += "?"  # 用问号替换无法识别的字符
    
    # 确保文本是有效的 UTF-8
    try:
        cleaned_text.encode('utf-8')
    except UnicodeEncodeError:
        # 如果仍然有问题，使用错误处理策略
        cleaned_text = cleaned_text.encode('utf-8', errors='replace').decode('utf-8')
    
    return cleaned_text

def is_not_empty(data):
    """
    安全地检查数据是否非空
    """
    if data is None:
        return False
    if isinstance(data, (list, tuple, np.ndarray)):
        return len(data) > 0
    if isinstance(data, dict):
        return len(data) > 0
    return bool(data)

def safe_get_data_from_result(result, key, index=0):
    """
    安全地从结果中获取数据
    """
    try:
        if key in result and is_not_empty(result[key]):
            data = result[key][index]
            # 处理 numpy 数组
            if isinstance(data, np.ndarray):
                if data.ndim == 0:  # 标量数组
                    return data.item()
                else:  # 多维数组
                    return data.tolist()
            return data
        return None
    except Exception as e:
        print(f"从结果中获取 {key} 数据时出错: {e}")
        return None

def safe_get_collection_data(collection, batch_size=50):
    """
    安全地获取集合数据，分批处理以避免编码问题
    """
    try:
        # 先尝试获取所有 ID
        all_ids = collection.get(include=[])["ids"]
        print(f"总共有 {len(all_ids)} 个文档")
        
        if len(all_ids) == 0:
            print("集合为空")
            return {"ids": [], "documents": [], "embeddings": [], "metadatas": []}
        
        all_results = {
            "ids": [],
            "documents": [],
            "embeddings": [],
            "metadatas": []
        }
        
        # 分批获取数据
        for i in range(0, len(all_ids), batch_size):
            batch_ids = all_ids[i:i + batch_size]
            try:
                batch_results = collection.get(
                    ids=batch_ids,
                    include=["documents", "embeddings", "metadatas"]
                )
                
                # 合并结果
                all_results["ids"].extend(batch_results["ids"])
                
                # 安全地处理文档数据
                if is_not_empty(batch_results["documents"]):
                    for doc in batch_results["documents"]:
                        # 处理 numpy 数组
                        if isinstance(doc, np.ndarray):
                            if doc.ndim == 0:  # 标量数组
                                all_results["documents"].append(str(doc.item()))
                            else:  # 多维数组
                                all_results["documents"].append(str(doc.tolist()))
                        else:
                            all_results["documents"].append(doc if doc is not None else "")
                else:
                    all_results["documents"].extend([None] * len(batch_results["ids"]))
                    
                # 安全地处理嵌入数据
                if is_not_empty(batch_results["embeddings"]):
                    for emb in batch_results["embeddings"]:
                        # 处理 numpy 数组
                        if isinstance(emb, np.ndarray):
                            if emb.ndim == 0:  # 标量数组
                                all_results["embeddings"].append(emb.item())
                            else:  # 多维数组
                                all_results["embeddings"].append(emb.tolist())
                        else:
                            all_results["embeddings"].append(emb)
                else:
                    all_results["embeddings"].extend([None] * len(batch_results["ids"]))
                    
                # 安全地处理元数据
                if is_not_empty(batch_results["metadatas"]):
                    all_results["metadatas"].extend(batch_results["metadatas"])
                else:
                    all_results["metadatas"].extend([None] * len(batch_results["ids"]))
                
                print(f"已处理 {min(i + batch_size, len(all_ids))}/{len(all_ids)} 个文档")
                
            except Exception as batch_error:
                print(f"处理批次 {i//batch_size + 1} 时出错: {batch_error}")
                # 尝试单个获取
                for single_id in batch_ids:
                    try:
                        single_result = collection.get(
                            ids=[single_id],
                            include=["documents", "embeddings", "metadatas"]
                        )
                        
                        # 使用安全方法提取数据
                        doc_data = safe_get_data_from_result(single_result, "documents")
                        emb_data = safe_get_data_from_result(single_result, "embeddings")
                        meta_data = safe_get_data_from_result(single_result, "metadatas")
                        
                        all_results["ids"].append(single_result["ids"][0])
                        all_results["documents"].append(doc_data)
                        all_results["embeddings"].append(emb_data)
                        all_results["metadatas"].append(meta_data)
                            
                    except Exception as single_error:
                        print(f"获取单个文档 {single_id} 时出错: {single_error}")
                        # 添加占位符
                        all_results["ids"].append(single_id)
                        all_results["documents"].append("[无法读取文档内容]")
                        all_results["embeddings"].append(None)
                        all_results["metadatas"].append(None)
                
        return all_results
        
    except Exception as e:
        print(f"获取集合数据时出错: {e}")
        import traceback
        traceback.print_exc()
        return {"ids": [], "documents": [], "embeddings": [], "metadatas": []}

def diagnose_collection(collection_name="knowledge_base", host="localhost", port=8002):
    """
    诊断集合问题
    """
    print(f"正在连接到 ChromaDB ({host}:{port})...")
    
    try:
        # 使用标准客户端
        chroma_client = HttpClient(host=host, port=port)
        print("连接成功")
        
        print(f"获取集合 '{collection_name}'...")
        collection = chroma_client.get_or_create_collection(name=collection_name)
        
        print(f"集合中的文档数量: {collection.count()}")
        
        if collection.count() == 0:
            print("集合为空，无需进一步处理")
            return
        
        print("尝试安全获取数据...")
        results = safe_get_collection_data(collection)
        
        print(f"成功获取 {len(results['ids'])} 个文档")
        
        # 显示前几个文档的内容
        displayed = 0
        for i in range(len(results['documents'])):
            if displayed >= 3:
                break
            doc_content = results['documents'][i]
            if doc_content and doc_content != "[无法读取文档内容]":
                # 只显示前100个字符
                if isinstance(doc_content, str):
                    preview = doc_content[:100] + "..." if len(doc_content) > 100 else doc_content
                    print(f"文档 {displayed+1} 预览: {preview}")
                    displayed += 1
                else:
                    print(f"文档 {displayed+1} 内容类型: {type(doc_content)}")
                    displayed += 1
        
        return results
        
    except Exception as e:
        print(f"诊断过程中出错: {e}")
        import traceback
        traceback.print_exc()
        return None

def main():
    """
    主函数
    """
    print("ChromaDB 编码问题诊断和修复工具")
    print("=" * 40)
    
    # 可以通过命令行参数指定集合名称和连接信息
    import argparse
    parser = argparse.ArgumentParser(description='诊断和修复 ChromaDB 编码问题')
    parser.add_argument('--collection', default='knowledge_base', help='集合名称')
    parser.add_argument('--host', default='localhost', help='ChromaDB 主机地址')
    parser.add_argument('--port', type=int, default=8002, help='ChromaDB 端口')
    
    args = parser.parse_args()
    
    results = diagnose_collection(
        collection_name=args.collection,
        host=args.host,
        port=args.port
    )
    
    if results:
        print("\n诊断完成!")
        print(f"成功处理了 {len(results['ids'])} 个文档")
    else:
        print("\n诊断失败，请检查错误信息")

if __name__ == "__main__":
    main()